package pl.edu.agh.som;

import pl.edu.agh.neural.core.ILayer;
import pl.edu.agh.neural.core.INeuron;
import pl.edu.agh.neural.core.ITrainableLayer;
import pl.edu.agh.neural.simple.InputConnection;
import pl.edu.agh.som.conscience.IConscienceFunction;
import pl.edu.agh.som.distance.IDistanceFunction;
import pl.edu.agh.som.learning.ILearningFunction;
import pl.edu.agh.som.neighborhood.INeighborhoodFunction;
import pl.edu.agh.som.winfrequency.IWinFrequencyUpdateFunction;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

public class KohonenMap implements ITrainableLayer
{

    private Random rand = new Random();
    private int dim;
    private Node[][] nodes;
    private int rows;
    private int columns;
    private IDistanceFunction distanceFunction;
    private ILearningFunction learningFunction;
    private INeighborhoodFunction neighborhoodFunction;
    private IConscienceFunction conscienceFunction;
    private IWinFrequencyUpdateFunction winFrequencyUpdateFunction;

    public KohonenMap(int dataVectorDim,
                      int rows,
                      int columns,
                      IDistanceFunction distanceFunction,
                      ILearningFunction learningFunction,
                      INeighborhoodFunction neighborhoodFunction,
                      IConscienceFunction conscienceFunction,
                      IWinFrequencyUpdateFunction winFrequencyUpdateFunction,
                      ILayer previousLayer)
    {
        this.dim = dataVectorDim;
        this.rows = rows;
        this.columns = columns;
        this.distanceFunction = distanceFunction;
        this.learningFunction = learningFunction;
        this.neighborhoodFunction = neighborhoodFunction;
        this.conscienceFunction = conscienceFunction;
        this.winFrequencyUpdateFunction = winFrequencyUpdateFunction;

        nodes = new Node[rows][];
        for (int i = 0; i < rows; i++)
        {
            nodes[i] = new Node[columns];
            for (int j = 0; j < columns; j++)
            {
                InputConnection[] inputConnections = createInputConnections(previousLayer);
                nodes[i][j] = new Node(i, j, 0.1, inputConnections);

                double[] currentWeights = nodes[i][j].getInputConnectionsWeights();
                double[] newNormalizedWeights = normalize(currentWeights);
                nodes[i][j].setInputConnectionsWeights(newNormalizedWeights);
            }
        }
    }

    @Override
    public List<? extends INeuron> getNeurons()
    {
        List<Node> nodesList = new ArrayList<>(rows * columns);
        for (int i = 0; i < rows; i++)
        {
            for (int j = 0; j < columns; j++)
            {
                nodesList.add(nodes[i][j]);
            }
        }

        return nodesList;
    }

    @Override
    public double[] evaluate()
    {
        Node winningNode = evaluateAndGetWinner();
        for (Node[] nodesRow : nodes)
        {
            for (Node node : nodesRow)
            {
                if (node != winningNode)
                {
                    node.setDistance(0);
                }
                else
                {
                    node.setDistance(1.0);
                }
            }
        }
        return winningNode.getInputConnectionsWeights();
    }

    public void train(int step)
    {
        Node winner = calculateDistancesAndFindWinner();

        for (Node[] nodesRow : nodes)
        {
            for (Node node : nodesRow)
            {
                double neighborhoodFactor = neighborhoodFunction.evaluate(node, winner, step);
                double learningFactor = learningFunction.evaluate(step);
                double[] currentWeights = node.getInputConnectionsWeights();
                double[] diff = vectorDiff(node.getInputConnectionsNodesValues(), currentWeights);
                double[] weightsDelta = vectorMul(neighborhoodFactor * learningFactor, diff);
                double[] newWeights = vectorSum(currentWeights, weightsDelta);
                node.setInputConnectionsWeights(newWeights);

                double newWinFrequency = winFrequencyUpdateFunction.evaluate(node, node.equals(winner));
                node.setWinFrequency(newWinFrequency);
            }
        }
    }

    public Node trainAndGetWinner(int step)
    {
        Node winner = calculateDistancesAndFindWinner();

        Node node = winner;
//                double neighborhoodFactor = neighborhoodFunction.evaluate(node, winner, step);
        double learningFactor = learningFunction.evaluate(step);
        double[] currentWeights = node.getInputConnectionsWeights();
        double[] diff = vectorDiff(node.getInputConnectionsNodesValues(), currentWeights);
        double[] weightsDelta = vectorMul(learningFactor, diff);
        double[] newWeights = vectorSum(currentWeights, weightsDelta);
        double[] newNormalizedWeights = normalize(newWeights);
        node.setInputConnectionsWeights(newNormalizedWeights);

        double newWinFrequency = winFrequencyUpdateFunction.evaluate(node, node.equals(winner));
        node.setWinFrequency(newWinFrequency);

        for (Node[] nodesRow : nodes)
        {
            for (Node nodei : nodesRow)
            {
                if (nodei != winner)
                {
                    nodei.setDistance(0);
                }
                else
                {
                    nodei.setDistance(1.0);
                }
            }
        }

        return winner;
    }


    private double[] normalize(double[] newWeights)
    {
        double sum = 0;
        double[] normalized = Arrays.copyOf(newWeights, newWeights.length);
        for (double i : newWeights)
        {
            sum += i;
        }
        for (int i = 0; i < normalized.length; i++)
        {
            normalized[i] = newWeights[i] / sum;
        }
        return normalized;
    }

    public Node evaluateAndGetWinner()
    {
        return calculateDistancesAndFindWinner();
    }

    private double randomDoubleFromRange(double min, double max)
    {
        return min + (max - min) * rand.nextDouble();
    }

    private InputConnection[] createInputConnections(ILayer previousLayer)
    {
        List<? extends INeuron> neurons = previousLayer.getNeurons();
        InputConnection[] inputConnections = new InputConnection[neurons.size()];
        for (int i = 0; i < neurons.size(); i++)
        {
            inputConnections[i] = new InputConnection(neurons.get(i), randomDoubleFromRange(0.0, 1.0));
        }
        return inputConnections;
    }

    private Node calculateDistancesAndFindWinner()
    {

        Node winner = null;
        double minDistance = Double.MAX_VALUE;
        for (Node[] nodesRow : nodes)
        {
            for (Node node : nodesRow)
            {
                double distance = distanceFunction.evaluate(node.getInputConnectionsWeights(), node.getInputConnectionsNodesValues());
                if (conscienceFunction != null)
                {
                    double conscience = conscienceFunction.evaluate(node);
                    distance += conscience;
                }
                node.setDistance(distance);
                if (distance < minDistance)
                {
                    minDistance = distance;
                    winner = node;
                }
            }
        }

        return winner;
    }

    private double[] vectorSum(double[] vector1, double[] vector2)
    {
        double[] result = new double[dim];
        for (int i = 0; i < dim; i++)
        {
            result[i] = vector1[i] + vector2[i];
        }

        return result;
    }

    private double[] vectorDiff(double[] vector1, double[] vector2)
    {
        double[] result = new double[dim];
        for (int i = 0; i < dim; i++)
        {
            result[i] = vector1[i] - vector2[i];
        }

        return result;
    }

    private double[] vectorMul(double factor, double[] vector)
    {
        double[] result = new double[dim];
        for (int i = 0; i < dim; i++)
        {
            result[i] = factor * vector[i];
        }

        return result;
    }


}
